# coding:utf8
import cv2
import base64
import json

def drawImage(img,nodeList):
    for node in nodeList:
        for r in node['leaf']:
            x,y,w,h = r['bbox']
            type = node['type']
            id = node['id']
            index = r['index']
            text = '%d-%d-%s' %(id,index,type)
            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
            # putText https://blog.csdn.net/qq_41273999/article/details/134597738
            cv2.putText(img, text, (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
    imgStr = cv2.imencode('.jpg', img)[1].tostring()
    imgBase64 = base64.b64encode(imgStr).decode()
    return imgBase64

def getNodeInfo(img):
    regionList = mergeRegion(img)
    regionList.sort(key=lambda x: x['bbox'][1])
    #tableRegion = {'nodeId':1,'type':'table','bbox': '','cell':'','leaf':[{'index':1,'bbox':''}]}
    nodeList = []
    for region in regionList:
        nodeLen = len(nodeList)
        if nodeLen == 0:
          node = {'id':0,'type':region['type'],'leaf': [{'bbox':region['bbox'],'cell':region['cell']}]}
          nodeList.append(node)
        else:
            lastNode = nodeList[nodeLen-1]
            #跟节点的第一个叶子比进行
            lastNodeLeaf = lastNode['leaf'][0]
            regionHeight = region['bbox'][1] + region['bbox'][3]
            lastNodeLeafHeight = lastNodeLeaf['bbox'][1] + lastNodeLeaf['bbox'][3]
            if  lastNodeLeafHeight - regionHeight > 0 or abs(lastNodeLeafHeight - regionHeight) < 3:
                leaf = {'bbox':region['bbox']}
                lastNode['leaf'].append(leaf)
                nodeList[nodeLen-1] = lastNode
            else:
                node = {'id':lastNode['id']+1,'type':region['type'],'leaf': [{'bbox':region['bbox'],'cell':region['cell']}]}
                nodeList.append(node)
    for node in nodeList:
        leaf = node['leaf']
        leaf.sort(key=lambda x: x['bbox'][0])
        index = 0
        for item in leaf:
            item['index'] = index
            index = index + 1
        node['leaf'] = leaf
    height,width,_ = img.shape
    imgBase64 = drawImage(img,nodeList)
    nodeInfo = {'height':height,'width':width,'node':nodeList,'base64':imgBase64}
    return nodeInfo

def mergeRegion(img):
    mergeRegionList = []
    allRegionList = getAllRegion(img)
    tableRegionList = getTableRegion(img)
    if len(tableRegionList) == 0:
        mergeRegionList = allRegionList
    else: 
      for allRegion in allRegionList:
        for tableRegion in tableRegionList:
            # tableRegion 在allRegion里面，即改region为table
            if (allRegion['bbox'][0] < tableRegion['bbox'][0]) and (allRegion['bbox'][1] < tableRegion['bbox'][1]) \
            and (allRegion['bbox'][0]+allRegion['bbox'][2] > tableRegion['bbox'][0]+tableRegion['bbox'][2]) and (allRegion['bbox'][1] + allRegion['bbox'][3] > tableRegion['bbox'][1] + tableRegion['bbox'][3]):
            #if (allRegion['bbox'][0] >= tableRegion['bbox'][0]) and (allRegion['bbox'][1] >= tableRegion['bbox'][1]) \
            #and (allRegion['bbox'][0]+allRegion['bbox'][2] <= tableRegion['bbox'][0]+tableRegion['bbox'][2]) and (allRegion['bbox'][1] + allRegion['bbox'][3] <= tableRegion['bbox'][1] + tableRegion['bbox'][3]):
              allRegion = tableRegion
        mergeRegionList.append(allRegion)
    return mergeRegionList

def getAllRegion(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
    # 框的范围更大，会导致相邻合并，即使得相邻合并
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (12,12))
    # 膨胀更严重，会导致相邻合并
    dilate = cv2.dilate(thresh, kernel, iterations=2)
    #cv2.imwrite("dilate.png", dilate)
    regionList = []
    # 1. 查找轮廓
    contours = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 2. 筛选那些面积小的
    contours = contours[0] if len(contours) == 2 else contours[1]
    for c in contours:
        #x，y是矩阵左上点的坐标 x、y是距离边缘的距离
        box = cv2.boundingRect(c)
        region = {'type':'block','bbox': box,'cell':''}
        regionList.append(region)
    return regionList

def compareTable(tuple):
      return tuple[2] * tuple[3] 


def recursive_append(regionList,newRegionList):
    # 为1的时候是图片
    if len(regionList) <= 3:
        return
    else:
        maxRegion = max(regionList, key=compareTable)
        currentRegion = []
        newRegion = []
        for region in regionList:
            if (region[0] >= maxRegion[0]) and (region[1] >= maxRegion[1]) \
            and (region[0]+region[2] <= maxRegion[0]+maxRegion[2]) and (region[1] + region[3] <= maxRegion[1] + maxRegion[3]):
                currentRegion.append(region)
                continue
            else:
                newRegion.append(region)
        newRegionList.append(currentRegion)
        recursive_append(newRegion,newRegionList)
    

def getTableRegion(img):
    tableRegionList = []
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    binary = cv2.adaptiveThreshold(~gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 5, -5)
    regionList = []
    # 1. 查找轮廓
    contours, _ = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # 2. 筛选那些面积小的
    for i in range(len(contours)):
        cnt = contours[i]
        area = cv2.contourArea(cnt) 
        if(area < 2000):
            continue
        # 左上→左下→右下→右上
        box = cv2.boundingRect(cnt)
        regionList.append(box)
    newRegionList = []
    recursive_append(regionList,newRegionList)
    for newRegion in newRegionList:
        box = max(newRegion, key=compareTable)
        cellList = []
        for item in newRegion:
            if item != box:
                cellList.append(item)
        tableRegion = {'type':'table','bbox': box,'cell':formatCell(cellList,box[2],box[3])}
        tableRegionList.append(tableRegion)
    return tableRegionList

def formatCell(cellList,width,height):
    #cellList 按y分组
    cellGroup = {}
    for cell in cellList:
        key = cell[1]
        if key in cellGroup:
            colList = cellGroup[key]
            colList.append(cell)
        else:
            colList = []
            colList.append(cell)
        cellGroup.update({key:colList})
    #按y分组时，合并相似分组
    newCellGroup = {}
    current = 0
    for index in sorted(cellGroup,reverse=True):
        key = index 
        value = cellGroup[index]
        if abs(index-current) < 5:
            print('merge')
            merge = []
            for item in value:
                merge.append(item)
            for item in cellGroup[current]:
                merge.append(item)
            newCellGroup.update({current:merge})
        else:
            newCellGroup.update({key:value})
        current = index
    #计算表格基础信息
    rows = len(newCellGroup)
    colValue = []
    for key in sorted(newCellGroup):
        value = newCellGroup[key]
        colValue.append(len(value))
    cols = max(colValue)
    minWidth = width/cols
    minHeight = height/rows
    print('rows:%s cols:%s' % (rows,cols))
    #根据表格基础信息生成表格
    tableCellList = []
    row = 0
    col = 0
    for key in sorted(newCellGroup):
        value = newCellGroup[key]
        value.sort(key=lambda x: x[0])
        for item in value:
            #python3 除法是小数
            rowspan = 1 if item[3]//minHeight == 0 else item[3]//minHeight
            colspan = 1 if item[2]//minWidth == 0 else item[2]//minWidth
            tableCell = {'bbox':item,'row':row,'col':col,'rowspan':rowspan,'colspan':colspan}
            tableCellList.append(tableCell)
            col = col + colspan
        row = row + 1
        col = 0
    return tableCellList
    
if __name__ == '__main__':
    # 读取文件
    imagePath = '/data/zwxu/layoutparse/test.jpg'
    img = cv2.imread(imagePath)
    height,width,_ = img.shape
    print('width:%s height:%s' % (width,height))
    nodeInfo = getNodeInfo(img)
    result = {'errCode': '0000','errMsg': '', 'result': {'data' : nodeInfo }}
    print(json.dumps(result))
    imgdata = base64.b64decode(nodeInfo['base64'])
    #将图片保存为文件
    with open("all.jpg",'wb') as f:
        f.write(imgdata)

